package it.uniroma3.dia.ia.textcategorization.categorizer.opennlp;

import it.uniroma3.dia.ia.textcategorization.model.Category;
import it.uniroma3.dia.ia.textcategorization.model.Text;
import it.uniroma3.dia.ia.textcategorization.texcat.TexCatConfiguration;

import java.io.BufferedReader;
import java.io.BufferedWriter;
import java.io.File;
import java.io.FileInputStream;
import java.io.FileOutputStream;
import java.io.IOException;
import java.io.InputStream;
import java.io.InputStreamReader;
import java.io.OutputStreamWriter;
import java.io.Writer;
import java.util.Map;
import opennlp.tools.doccat.DocumentSample;
import opennlp.tools.doccat.DocumentSampleStream;
import opennlp.tools.util.ObjectStream;
import opennlp.tools.util.PlainTextByLineStream;

final class OpenNLPTrainingFormat {

	private static final String training_file_pathname = TexCatConfiguration
			.getWorkspacePathname() + "/opennlp/training.train";

	static ObjectStream<DocumentSample> getFormattedTrainingMaterial(
			Map<Text, Category> trainingSet) throws IOException {
		// Writing the formatted training material in a file
		ensureOpenNLPDirectoryExistence();
		FileOutputStream fileOutputStream = new FileOutputStream(
				training_file_pathname);
		BufferedWriter writer = new BufferedWriter(new OutputStreamWriter(
				fileOutputStream));
		for (Text text : trainingSet.keySet()) {
			Category category = trainingSet.get(text);
			writeCategory(writer, category);
			writeText(writer, text);
			writer.write("\n");
		}
		writer.close();

		// Giving the training material back as ObjectStream<DocumentSample>
		InputStream inputStream = new FileInputStream(training_file_pathname);
		ObjectStream<String> lineStream = new PlainTextByLineStream(
				inputStream, "UTF-8");
		ObjectStream<DocumentSample> trainingMaterial = new DocumentSampleStream(
				lineStream);
		return trainingMaterial;
	}

	private static void writeCategory(Writer writer, Category category)
			throws IOException {
		writer.write(category.asPath() + " ");
	}

	private static void writeText(Writer writer, Text text) throws IOException {
		BufferedReader reader = new BufferedReader(new InputStreamReader(
				text.asInputStream()));
		String line = reader.readLine();
		while (line != null) {
			writer.write(line + " ");
			line = reader.readLine();
		}
	}

	private static void ensureOpenNLPDirectoryExistence() {
		File working_directory = new File(
				TexCatConfiguration.getWorkspacePathname() + "/opennlp");
		if (!working_directory.exists()) {
			working_directory.mkdir();
		}
	}

}
